import torch
import torch.nn as nn


class DeepNeuralNetworkModel(nn.Module):

    def __init__(self, input_dim, output_dim):
        super(DeepNeuralNetworkModel, self).__init__()

        # define hyperparameters
        layer_input_size = 16
        layer_list = ['D', 32, 64, 128, 'D', 256, 128, 'D', 64, 32, 16]

        # initialize layers list
        self.layers = nn.ModuleList()

        # build model
        self.layers.append(nn.Linear(input_dim, layer_input_size))
        self.layers.append(nn.Tanh())

        for v in layer_list:
            if v == 'D':
                self.layers.append(nn.Dropout(0.01))
            else:
                self.layers.append(nn.Linear(layer_input_size, v))
                self.layers.append(nn.Tanh())
                self.layers.append(nn.BatchNorm1d(v))
                layer_input_size = v

        self.layers.append(nn.Linear(layer_input_size, output_dim))
        self.layers.append(nn.Sigmoid())

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x
